package com.evaluation_system.Util;
import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.GeneralSecurityException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

public class PdfSM4Encryptor {

        private static final String ALGORITHM = "SM4";
        private static final String TRANSFORMATION = "SM4/CBC/PKCS5Padding";
        private static final int KEY_SIZE = 128; // SM4固定使用128位密钥
        private static final int IV_SIZE = 16; // SM4块大小固定为16字节

        // ------------------- 加密相关 -------------------
        /**
         * 生成 SM4 密钥 (128位)
         */
        public static SecretKey generateKey() throws NoSuchAlgorithmException {
            KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
            keyGen.init(KEY_SIZE);
            return keyGen.generateKey();
        }

        /**
         * 生成随机初始化向量 (IV)
         */
        public static IvParameterSpec generateIv() {
            byte[] iv = new byte[IV_SIZE];
            new SecureRandom().nextBytes(iv);
            return new IvParameterSpec(iv);
        }

        /**
         * 加密 PDF 字节数据
         * @param pdfBytes 原始PDF字节数组
         * @param key      SM4密钥
         * @param iv       初始化向量
         * @return 加密后的字节数组
         */
        public static byte[] encrypt(byte[] pdfBytes, SecretKey key, IvParameterSpec iv)
                throws GeneralSecurityException {
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.ENCRYPT_MODE, key, iv);
            return cipher.doFinal(pdfBytes);
        }

        // ------------------- 解密相关 -------------------
        /**
         * 解密 PDF 字节数据
         * @param encryptedBytes 加密后的字节数组
         * @param key           SM4密钥 (必须与加密时相同)
         * @param iv            初始化向量 (必须与加密时相同)
         * @return 解密后的原始PDF字节数组
         */
        public static byte[] decrypt(byte[] encryptedBytes, SecretKey key, IvParameterSpec iv)
                throws GeneralSecurityException {
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.DECRYPT_MODE, key, iv);
            return cipher.doFinal(encryptedBytes);
        }

        // ------------------- 密钥转换工具 -------------------
        /**
         * 将密钥转换为Base64字符串 (便于存储或传输)
         */
        public static String keyToString(SecretKey key) {
            return Base64.getEncoder().encodeToString(key.getEncoded());
        }

        /**
         * 从Base64字符串恢复密钥
         */
        public static SecretKey stringToKey(String keyStr) {
            byte[] decodedKey = Base64.getDecoder().decode(keyStr);
            return new SecretKeySpec(decodedKey, ALGORITHM);
        }

        /**
         * 将IV转换为Base64字符串
         */
        public static String ivToString(IvParameterSpec iv) {
            return Base64.getEncoder().encodeToString(iv.getIV());
        }

        /**
         * 从Base64字符串恢复IV
         */
        public static IvParameterSpec stringToIv(String ivStr) {
            byte[] ivBytes = Base64.getDecoder().decode(ivStr);
            return new IvParameterSpec(ivBytes);
        }

}
